from asyncio import FastChildWatcher
from code import interact
import os
import time
import numpy as np
from isaacgym import gymutil
from isaacgym import gymapi
#from isaacgym import gymtorch
from math import sqrt
import math
from sympy import false
import cv2
from draw import *
from pcgworker.PCGWorker import *

from wfc_vecenv_stable_baselines import *

from stable_baselines3 import PPO
from stable_baselines3.common.callbacks import BaseCallback
from stable_baselines3.common.results_plotter import load_results, ts2xy
from stable_baselines3.common.monitor import Monitor
from stable_baselines3.common import results_plotter
from stable_baselines3.common.results_plotter import load_results, ts2xy, plot_results
from stable_baselines3.common.noise import NormalActionNoise
from stable_baselines3.common.callbacks import BaseCallback

import torch


LOGDIR = "./training_logs"

m_env = PCGVecEnvStableBaselines(seed_pth = "pcgworker/new.json", headless_ = False)
# m_env = PCGVecEnvStableBaselines(headless_ = False)
m_env.reset()

print("num_envs:",m_env.num_envs)

all_obs = np.zeros((16, 84, 84, 4), dtype=np.uint8)

all_actions = np.zeros(m_env.num_matrix_envs)
for i in range(m_env.num_matrix_envs):
    all_actions[i] = -2

steps_ = 0
reset_id = 0

PCGWorker_ = PCGWorker(9,9)

manual = False

m_env.pause()

while True:

    # steps_ += 1

    observation, reward, done, info = m_env.step(all_actions)

    obs1 = observation[0]
    # transpose obs1 from (3,84,84) to (84,84,3)
    obs1 = np.transpose(obs1, (1,2,0))

    # create a black opencv image of size(84,84)
    img = np.zeros((84,84,3), np.uint8)
    # display img
    cv2.imshow('img', obs1)
    key_ = cv2.waitKey(1)
    if key_ == 119:  # w
        action = 0
    elif key_ == 115:    # s
        action = 1
    elif key_ == 97:     # a
        action = 2
    elif key_ == 100:    # d
        action = 3
    elif key_ == 107:    # ccw
        action = 4
    elif key_ == 108:    # cw
        action = 5
    elif key_ == 105:    # reset
        m_env.reset()
    elif key_ == 114:    # resume
        m_env.resume()
        manual = True
    elif key_ == 112:    # pause
        m_env.pause()
    elif key_ == 109:    # change landscape
        seed = PCGWorker_.generate()
        m_env.set_landscape(env_id = 0,seed_ = seed)
    else:
        action = -2

    # print(key_)

    for i in range(m_env.num_matrix_envs):
        all_actions[i] = action

    # # if manual:
    # #     key_ = cv2.waitKey(0)
    # #     if key_ == 109:
    # #         manual = False

    m_env.render()
